package git.soulbgm.common.advice;

import com.fasterxml.jackson.databind.ObjectMapper;
import git.soulbgm.common.exception.EncryptException;
import git.soulbgm.config.KeyConfig;
import git.soulbgm.pojo.RsaVo;
import git.soulbgm.utils.AesUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.MethodParameter;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import static git.soulbgm.config.KeyConfig.HEADER_CK;

/**
 * 加密响应
 *
 * @author SoulBGM
 * @date 2023-10-08
 */
@ControllerAdvice(basePackages = {"git.soulbgm.controller"})
public class EncryptResponseAdvice implements ResponseBodyAdvice<Object> {

    private static final Logger log = LoggerFactory.getLogger(EncryptResponseAdvice.class);
    private final boolean encrypt;
    private final String uniqueKey;
    private final ObjectMapper mapper;

    public EncryptResponseAdvice(ObjectMapper mapper, KeyConfig config) {
        this.encrypt = config.isEncrypt();
        this.uniqueKey = config.getUniqueKey();
        this.mapper = mapper;
    }

    @Override
    public boolean supports(MethodParameter methodParameter, Class aClass) {
        return encrypt;
    }

    /**
     * 返回结果加密
     *
     * @param body            接口返回的对象
     * @param methodParameter 方法参数
     * @param mediaType       媒体类型
     * @param aClass          Http消息转换器Class
     * @param request         请求
     * @param response        响应
     * @return obj
     */
    @Override
    public Object beforeBodyWrite(Object body, MethodParameter methodParameter, MediaType mediaType, Class aClass, ServerHttpRequest request, ServerHttpResponse response) {
        String ck = request.getHeaders().getFirst(HEADER_CK);
        if (uniqueKey.equals(ck)) {
            return body;
        }
        if (encrypt) {
            try {
                String content = mapper.writer().writeValueAsString(body);
                String result = AesUtil.encrypt(content, uniqueKey);
                return RsaVo.get(result);
            } catch (Exception e) {
                log.error("Encrypted data exception", e);
                throw new EncryptException("加密失败");
            }
        }
        return body;
    }

}